#include "diskid.h"


// Define global buffers.
#define PRINTING_TO_CONSOLE_ALLOWED


/*==============================================================================*/
DiskId::DiskId()
{

}


long DiskId::getDiskID()
{
    int done = FALSE;
    // char string [1024];
    __int64 id = 0;
    OSVERSIONINFO version;

    strcpy(HardDriveSerialNumber, "");

    memset(&version, 0, sizeof (version));
    version.dwOSVersionInfoSize = sizeof (OSVERSIONINFO);
    GetVersionEx(&version);
    if (version.dwPlatformId == VER_PLATFORM_WIN32_NT)
    {
        done = ReadPhysicalDriveInNTWithAdminRights();

        //  this should work in WinNT or Win2K if previous did not work
        //  this is kind of a backdoor via the SCSI mini port driver into
        //     the IDE drives
        done = ReadIdeDriveAsScsiDriveInNT();

        //  this works under WinNT4 or Win2K or WinXP if you have any rights
        done = ReadPhysicalDriveInNTWithZeroRights();

        //  this works under WinNT4 or Win2K or WinXP or Windows Server 2003 or Vista if you have any rights
        done = ReadPhysicalDriveInNTUsingSmart();
    }
    else
    {
        //  this works under Win9X and calls a VXD
        int attempt = 0;

        //  try this up to 10 times to get a hard drive serial number
        for (attempt = 0;
             attempt < 10 && !done && 0 == HardDriveSerialNumber[0];
             attempt++)
            done = ReadDrivePortsInWin9X();
    }

    if (HardDriveSerialNumber[0] > 0)
    {
        char *p = HardDriveSerialNumber;

        //  ignore first 5 characters from western digital hard drives if
        //  the first four characters are WD-W
        if (!strncmp(HardDriveSerialNumber, "WD-W", 4))
            p += 5;
        for (; p && *p; p++)
        {
            if ('-' == *p)
                continue;
            id *= 10;
            switch (*p)
            {
            case '0': id += 0; break;
            case '1': id += 1; break;
            case '2': id += 2; break;
            case '3': id += 3; break;
            case '4': id += 4; break;
            case '5': id += 5; break;
            case '6': id += 6; break;
            case '7': id += 7; break;
            case '8': id += 8; break;
            case '9': id += 9; break;
            case 'a': case 'A': id += 10; break;
            case 'b': case 'B': id += 11; break;
            case 'c': case 'C': id += 12; break;
            case 'd': case 'D': id += 13; break;
            case 'e': case 'E': id += 14; break;
            case 'f': case 'F': id += 15; break;
            case 'g': case 'G': id += 16; break;
            case 'h': case 'H': id += 17; break;
            case 'i': case 'I': id += 18; break;
            case 'j': case 'J': id += 19; break;
            case 'k': case 'K': id += 20; break;
            case 'l': case 'L': id += 21; break;
            case 'm': case 'M': id += 22; break;
            case 'n': case 'N': id += 23; break;
            case 'o': case 'O': id += 24; break;
            case 'p': case 'P': id += 25; break;
            case 'q': case 'Q': id += 26; break;
            case 'r': case 'R': id += 27; break;
            case 's': case 'S': id += 28; break;
            case 't': case 'T': id += 29; break;
            case 'u': case 'U': id += 30; break;
            case 'v': case 'V': id += 31; break;
            case 'w': case 'W': id += 32; break;
            case 'x': case 'X': id += 33; break;
            case 'y': case 'Y': id += 34; break;
            case 'z': case 'Z': id += 35; break;
            }
        }
    }

    id %= 100000000;
    if (strstr(HardDriveModelNumber, "IBM-"))
        id += 300000000;
    else if (strstr(HardDriveModelNumber, "MAXTOR") ||
             strstr(HardDriveModelNumber, "Maxtor"))
        id += 400000000;
    else if (strstr(HardDriveModelNumber, "WDC "))
        id += 500000000;
    else
        id += 600000000;

#ifdef PRINTING_TO_CONSOLE_ALLOWED

    printf("\nHard Drive Serial Number__________: %s\n",
           HardDriveSerialNumber);
    printf("\nHard Drive Model Number___________: %s\n",
           HardDriveModelNumber);
    printf("\nComputer ID_______________________: %lld\n", id);

#endif

    return static_cast<long>(id);
}



int DiskId::ReadPhysicalDriveInNTWithAdminRights()
{
    int done = FALSE;
    int drive = 0;

    for (drive = 0; drive < MAX_IDE_DRIVES; drive++)
    {
        HANDLE hPhysicalDriveIOCTL = nullptr;

        //  Try to get a handle to PhysicalDrive IOCTL, report failure
        //  and exit if can't.
        WCHAR driveName[256];

        swprintf(driveName, L"\\\\.\\PhysicalDrive%d", drive);

        //  Windows NT, Windows 2000, must have admin rights
        hPhysicalDriveIOCTL = CreateFile(driveName,
                                         GENERIC_READ | GENERIC_WRITE,
                                         FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                                         OPEN_EXISTING, 0, nullptr);

        if (hPhysicalDriveIOCTL != INVALID_HANDLE_VALUE)
        {
            GETVERSIONOUTPARAMS VersionParams;
            DWORD               cbBytesReturned = 0;

            // Get the version, etc of PhysicalDrive IOCTL
            memset((void*)&VersionParams, 0, sizeof(VersionParams));

            DeviceIoControl(hPhysicalDriveIOCTL, DFP_GET_VERSION,
                                 nullptr,
                                 0,
                                 &VersionParams,
                                 sizeof(VersionParams),
                                 &cbBytesReturned, nullptr);

            // If there is a IDE device at number "i" issue commands
            // to the device
            if (VersionParams.bIDEDeviceMap > 0)
            {
                BYTE             bIDCmd = 0;   // IDE or ATAPI IDENTIFY cmd
                SENDCMDINPARAMS  scip;
                //SENDCMDOUTPARAMS OutCmd;

                // Now, get the ID sector for all IDE devices in the system.
                // If the device is ATAPI use the IDE_ATAPI_IDENTIFY command,
                // otherwise use the IDE_ATA_IDENTIFY command
                bIDCmd = (VersionParams.bIDEDeviceMap >> drive & 0x10) ? \
                            IDE_ATAPI_IDENTIFY : IDE_ATA_IDENTIFY;

                memset(&scip, 0, sizeof(scip));
                memset(IdOutCmd, 0, sizeof(IdOutCmd));

                if (DoIDENTIFY(hPhysicalDriveIOCTL,
                               &scip,
                               (PSENDCMDOUTPARAMS)&IdOutCmd,
                               static_cast<BYTE>(bIDCmd),
                               static_cast<BYTE>(drive),
                               &cbBytesReturned))
                {
                    DWORD diskdata[256];
                    int ijk = 0;
                    USHORT *pIdSector = (USHORT *)
                            ((PSENDCMDOUTPARAMS)IdOutCmd)->bBuffer;

                    for (ijk = 0; ijk < 256; ijk++)
                        diskdata[ijk] = pIdSector[ijk];

                    //PrintIdeInfo(drive, diskdata);

                    done = TRUE;
                }
            }

            CloseHandle(hPhysicalDriveIOCTL);
        }
    }

    return done;
}

int DiskId::ReadIdeDriveAsScsiDriveInNT()
{
    int done = FALSE;
    int controller = 0;

    for (controller = 0; controller < 16; controller++)
    {
        HANDLE hScsiDriveIOCTL = nullptr;
        WCHAR   driveName[256];

        //  Try to get a handle to PhysicalDrive IOCTL, report failure
        //  and exit if can't.
        swprintf(driveName, L"\\\\.\\Scsi%d:", controller);

        //  Windows NT, Windows 2000, any rights should do
        hScsiDriveIOCTL = CreateFile(driveName,
                                     GENERIC_READ | GENERIC_WRITE,
                                     FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                                     OPEN_EXISTING, 0, nullptr);

        if (hScsiDriveIOCTL != INVALID_HANDLE_VALUE)
        {
            int drive = 0;

            for (drive = 0; drive < 2; drive++)
            {
                char buffer[sizeof (SRB_IO_CONTROL)+SENDIDLENGTH];
                SRB_IO_CONTROL *p = (SRB_IO_CONTROL *)buffer;
                SENDCMDINPARAMS *pin =
                        (SENDCMDINPARAMS *)(buffer + sizeof (SRB_IO_CONTROL));
                DWORD dummy;

                memset(buffer, 0, sizeof (buffer));
                p->HeaderLength = sizeof (SRB_IO_CONTROL);
                p->Timeout = 10000;
                p->Length = SENDIDLENGTH;
                p->ControlCode = IOCTL_SCSI_MINIPORT_IDENTIFY;
                strncpy((char *)p->Signature, "SCSIDISK", 8);

                pin->irDriveRegs.bCommandReg = IDE_ATA_IDENTIFY;
                pin->bDriveNumber = drive;

                if (DeviceIoControl(hScsiDriveIOCTL, IOCTL_SCSI_MINIPORT,
                                    buffer,
                                    sizeof (SRB_IO_CONTROL)+
                                    sizeof (SENDCMDINPARAMS)-1,
                                    buffer,
                                    sizeof (SRB_IO_CONTROL)+SENDIDLENGTH,
                                    &dummy, nullptr))
                {
                    SENDCMDOUTPARAMS *pOut =
                            (SENDCMDOUTPARAMS *)(buffer + sizeof (SRB_IO_CONTROL));
                    IDSECTOR *pId = (IDSECTOR *)(pOut->bBuffer);
                    if (pId->sModelNumber[0])
                    {
                        DWORD diskdata[256];
                        int ijk = 0;
                        USHORT *pIdSector = (USHORT *)pId;

                        for (ijk = 0; ijk < 256; ijk++)
                            diskdata[ijk] = pIdSector[ijk];

                        //PrintIdeInfo(controller * 2 + drive, diskdata);

                        done = TRUE;
                    }
                }
            }
            CloseHandle(hScsiDriveIOCTL);
        }
    }

    return done;
}

int DiskId::ReadPhysicalDriveInNTWithZeroRights()
{
    int done = FALSE;
    int drive = 0;

    for (drive = 0; drive < MAX_IDE_DRIVES; drive++)
    {
        HANDLE hPhysicalDriveIOCTL = nullptr;

        //  Try to get a handle to PhysicalDrive IOCTL, report failure
        //  and exit if can't.
        WCHAR driveName[256];

        swprintf(driveName, L"\\\\.\\PhysicalDrive%d", drive);

        //  Windows NT, Windows 2000, Windows XP - admin rights not required
        hPhysicalDriveIOCTL = CreateFile((driveName), 0,
                                         FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                                         OPEN_EXISTING, 0, nullptr);
        if (hPhysicalDriveIOCTL != INVALID_HANDLE_VALUE)
        {
            STORAGE_PROPERTY_QUERY query;
            DWORD cbBytesReturned = 0;
            char local_buffer[10000];

            memset((void *)& query, 0, sizeof (query));
            query.PropertyId = StorageDeviceProperty;
            query.QueryType = PropertyStandardQuery;

            memset(local_buffer, 0, sizeof (local_buffer));

            if (DeviceIoControl(hPhysicalDriveIOCTL, IOCTL_STORAGE_QUERY_PROPERTY,
                                &query,
                                sizeof (query),
                                &local_buffer[0],
                                sizeof (local_buffer),
                                &cbBytesReturned, nullptr))
            {
                STORAGE_DEVICE_DESCRIPTOR * descrip = (STORAGE_DEVICE_DESCRIPTOR *)& local_buffer;
                char serialNumber[1000];
                char modelNumber[1000];
//                char vendorId[1000];
                char productRevision[1000];

//                flipAndCodeBytes(local_buffer,
//                                 static_cast<int>(descrip->VendorIdOffset),
//                                 0, vendorId);
                flipAndCodeBytes(local_buffer,
                                 descrip->ProductIdOffset,
                                 0, modelNumber);
                flipAndCodeBytes(local_buffer,
                                 descrip->ProductRevisionOffset,
                                 0, productRevision);
                flipAndCodeBytes(local_buffer,
                                 descrip->SerialNumberOffset,
                                 1, serialNumber);

                if (0 == HardDriveSerialNumber[0] &&
                        //  serial number must be alphanumeric
                        //  (but there can be leading spaces on IBM drives)
                        (isalnum(serialNumber[0]) || isalnum(serialNumber[19])))
                {
                    strcpy(HardDriveSerialNumber, serialNumber);
                    strcpy(HardDriveModelNumber, modelNumber);
                    done = TRUE;
                }
#ifdef PRINTING_TO_CONSOLE_ALLOWED
                printf("\n**** STORAGE_DEVICE_DESCRIPTOR for drive %d ****\n"
//                       "Vendor Id = [%s]\n"
                       "Product Id = [%s]\n"
                       "Product Revision = [%s]\n"
                       "Serial Number = [%s]\n",
                       drive,
//                       vendorId,
                       modelNumber,
                       productRevision,
                       serialNumber);
#endif
            }
            else
            {
                DWORD err = GetLastError();
            }

            CloseHandle(hPhysicalDriveIOCTL);
        }
    }

    return done;
}

int DiskId::ReadPhysicalDriveInNTUsingSmart()
{
    int done = FALSE;
    int drive = 0;

    for (drive = 0; drive < MAX_IDE_DRIVES; drive++)
    {
        HANDLE hPhysicalDriveIOCTL = nullptr;

        //  Try to get a handle to PhysicalDrive IOCTL, report failure
        //  and exit if can't.
        WCHAR driveName[256];

        swprintf(driveName, L"\\\\.\\PhysicalDrive%d", drive);

        //  Windows NT, Windows 2000, Windows Server 2003, Vista
        hPhysicalDriveIOCTL = CreateFile((driveName),
                                         GENERIC_READ | GENERIC_WRITE,
                                         FILE_SHARE_DELETE | FILE_SHARE_READ | FILE_SHARE_WRITE,
                                         nullptr, OPEN_EXISTING, 0, nullptr);

        if (hPhysicalDriveIOCTL == INVALID_HANDLE_VALUE)
        {
            //错误
        }
        else
        {
            GETVERSIONINPARAMS GetVersionParams;
            DWORD cbBytesReturned = 0;

            // Get the version, etc of PhysicalDrive IOCTL
            memset((void*)& GetVersionParams, 0, sizeof(GetVersionParams));

            if (!DeviceIoControl(hPhysicalDriveIOCTL, SMART_GET_VERSION,
                                 nullptr,
                                 0,
                                 &GetVersionParams, sizeof (GETVERSIONINPARAMS),
                                 &cbBytesReturned, nullptr))
            {

            }
            else
            {
                // Print the SMART version
                // PrintVersion (& GetVersionParams);
                // Allocate the command buffer
                ULONG CommandSize = sizeof(SENDCMDINPARAMS)+IDENTIFY_BUFFER_SIZE;
                PSENDCMDINPARAMS Command = (PSENDCMDINPARAMS)malloc(CommandSize);
                // Retrieve the IDENTIFY data
                // Prepare the command

                Command->irDriveRegs.bCommandReg = ID_CMD;
                DWORD BytesReturned = 0;
                if (DeviceIoControl(hPhysicalDriveIOCTL,
                                    SMART_RCV_DRIVE_DATA, Command, sizeof(SENDCMDINPARAMS),
                                    Command, CommandSize,
                                    &BytesReturned, nullptr))
                {
                    done = TRUE;
                }

                // Done
                CloseHandle(hPhysicalDriveIOCTL);
                free(Command);
            }
        }
    }

    return done;
}

int DiskId::ReadDrivePortsInWin9X()
{
    int done = FALSE;
    unsigned long int i = 0;

    HANDLE VxDHandle = nullptr;
    pt_IdeDInfo pOutBufVxD = nullptr;
    DWORD lpBytesReturned = 0;

    //  set the thread priority high so that we get exclusive access to the disk
    BOOL status =
            // SetThreadPriority (GetCurrentThread(), THREAD_PRIORITY_TIME_CRITICAL);
            SetPriorityClass(GetCurrentProcess(), REALTIME_PRIORITY_CLASS);
    // SetPriorityClass (GetCurrentProcess (), HIGH_PRIORITY_CLASS);

    // 1. Make an output buffer for the VxD
    rt_IdeDInfo info;
    pOutBufVxD = &info;

    // KLUDGE WARNING!!!
    // HAVE to zero out the buffer space for the IDE information!
    // If this is NOT done then garbage could be in the memory
    // locations indicating if a disk exists or not.
    ZeroMemory(&info, sizeof(info));

    // 1. Try to load the VxD
    //  must use the short file name path to open a VXD file
    VxDHandle = CreateFile(L"\\\\.\\IDE21201.VXD", 0, 0, nullptr,
                           0, FILE_FLAG_DELETE_ON_CLOSE, nullptr);

    if (VxDHandle != INVALID_HANDLE_VALUE)
    {
        // 2. Run VxD function
        DeviceIoControl(VxDHandle, m_cVxDFunctionIdesDInfo,
                        nullptr, 0, pOutBufVxD, sizeof(pt_IdeDInfo), &lpBytesReturned, nullptr);

        // 3. Unload VxD
        CloseHandle(VxDHandle);
    }

    // 4. Translate and store data
    for (i = 0; i<8; i++)
    {
        if ((pOutBufVxD->DiskExists[i]) && (pOutBufVxD->IDEExists[i / 2]))
        {
            DWORD diskinfo[256];
            for (int j = 0; j < 256; j++)
                diskinfo[j] = pOutBufVxD->DisksRawInfo[i * 256 + j];

            // process the information for this buffer
            //PrintIdeInfo(i, diskinfo);
            done = TRUE;
        }
    }

    //  reset the thread priority back to normal
    // SetThreadPriority (GetCurrentThread(), THREAD_PRIORITY_NORMAL);
    SetPriorityClass(GetCurrentProcess(), NORMAL_PRIORITY_CLASS);

    return done;
}

WINBOOL DiskId::DoIDENTIFY(HANDLE hPhysicalDriveIOCTL, PSENDCMDINPARAMS pSCIP, PSENDCMDOUTPARAMS pSCOP, BYTE bIDCmd, BYTE bDriveNum, PDWORD lpcbBytesReturned)
{
    // Set up data structures for IDENTIFY command.
    pSCIP->cBufferSize = IDENTIFY_BUFFER_SIZE;
    pSCIP->irDriveRegs.bFeaturesReg = 0;
    pSCIP->irDriveRegs.bSectorCountReg = 1;
    //pSCIP -> irDriveRegs.bSectorNumberReg = 1;
    pSCIP->irDriveRegs.bCylLowReg = 0;
    pSCIP->irDriveRegs.bCylHighReg = 0;

    // Compute the drive number.
    pSCIP->irDriveRegs.bDriveHeadReg = 0xA0 | ((bDriveNum & 1) << 4);

    // The command can either be IDE identify or ATAPI identify.
    pSCIP->irDriveRegs.bCommandReg = bIDCmd;
    pSCIP->bDriveNumber = bDriveNum;
    pSCIP->cBufferSize = IDENTIFY_BUFFER_SIZE;

    return (DeviceIoControl(hPhysicalDriveIOCTL, DFP_RECEIVE_DRIVE_DATA,
                            static_cast<LPVOID>(pSCIP),
                            sizeof(SENDCMDINPARAMS)-1,
                            static_cast<LPVOID>(pSCOP),
                            sizeof(SENDCMDOUTPARAMS)+IDENTIFY_BUFFER_SIZE - 1,
                            lpcbBytesReturned, nullptr));
}

char *DiskId::flipAndCodeBytes(const char *str, int pos, int flip, char *buf)
{
    int i;
    int j = 0;
    int k = 0;

    buf[0] = '\0';
    if (pos <= 0)
        return buf;

    if (!j)
    {
        char p = 0;

        // First try to gather all characters representing hex digits only.
        j = 1;
        k = 0;
        buf[k] = 0;
        for (i = pos; j && str[i] != '\0'; ++i)
        {
            char c = static_cast<char>(tolower(str[i]));

            if (isspace(c))
                c = '0';

            ++p;
            buf[k] <<= 4;

            if (c >= '0' && c <= '9')
                buf[k] |= static_cast<unsigned char>(c - '0');
            else if (c >= 'a' && c <= 'f')
                buf[k] |= static_cast<unsigned char>(c - 'a' + 10);
            else
            {
                j = 0;
                break;
            }

            if (p == 2)
            {
                if (buf[k] != '\0' && !isprint(buf[k]))
                {
                    j = 0;
                    break;
                }
                ++k;
                p = 0;
                buf[k] = 0;
            }
        }
    }

    if (!j)
    {
        // There are non-digit characters, gather them as is.
        j = 1;
        k = 0;
        for (i = pos; j && str[i] != '\0'; ++i)
        {
            char c = str[i];

            if (!isprint(c))
            {
                j = 0;
                break;
            }

            buf[k++] = c;
        }
    }

    if (!j)
    {
        // The characters are not there or are not printable.
        k = 0;
    }

    buf[k] = '\0';

    if (flip)
        // Flip adjacent characters
        for (j = 0; j < k; j += 2)
        {
            char t = buf[j];
            buf[j] = buf[j + 1];
            buf[j + 1] = t;
        }

    // Trim any beginning and end space
    i = j = -1;
    for (k = 0; buf[k] != '\0'; ++k)
    {
        if (!isspace(buf[k]))
        {
            if (i < 0)
                i = k;
            j = k;
        }
    }

    if ((i >= 0) && (j >= 0))
    {
        for (k = i; (k <= j) && (buf[k] != '\0'); ++k)
            buf[k - i] = buf[k];
        buf[k - i] = '\0';
    }

    return buf;
}
